package com.web.crawler.compare;

import java.util.*;
import java.util.stream.Collectors;

public class Vectorizer {
    // 构建词袋模型
    public static Map<String, Integer> buildVocabulary(List<String> documents) {
        Map<String, Integer> vocabulary = new HashMap<>();
        int index = 0;

        // 遍历每个文档，构建词袋模型
        for (String doc : documents) {
            String[] words = doc.toLowerCase().split("\\s+");
            for (String word : words) {
                //  如果词不在词袋模型中，则添加
                if (!vocabulary.containsKey(word)) {
                    vocabulary.put(word, index++);
                }
            }
        }

        return vocabulary;
    }

    // 将文档转换为TF向量
    public static double[] toTfVector(String document, Map<String, Integer> vocabulary) {
        double[] vector = new double[vocabulary.size()];
        String[] words = document.toLowerCase().split("\\s+");

        for (String word : words) {
            Integer index = vocabulary.get(word);
            if (index != null) {
                vector[index]++;
            }
        }

        return vector;
    }

    // 计算TF-IDF向量
    public static double[] toTfIdfVector(String document, Map<String, Integer> vocabulary,
                                         Map<String, Double> idfMap) {
        double[] tfVector = toTfVector(document, vocabulary);
        double[] tfIdfVector = new double[tfVector.length];

        // 遍历TF向量，计算TF-IDF值
        for (Map.Entry<String, Integer> entry : vocabulary.entrySet()) {
            String word = entry.getKey();
            int index = entry.getValue();
            double idf = idfMap.getOrDefault(word, 1.0);
            tfIdfVector[index] = tfVector[index] * idf;
        }

        return tfIdfVector;
    }

    // 计算IDF值
    public static Map<String, Double> calculateIdf(List<String> documents,
                                                   Map<String, Integer> vocabulary) {
        Map<String, Double> idfMap = new HashMap<>();
        int totalDocs = documents.size();

        for (String word : vocabulary.keySet()) {
            int docCount = 0; // 统计包含当前单词的文档数量
            for (String doc : documents) {
                if (doc.toLowerCase().contains(word)) {
                    docCount++;
                }
            }
            idfMap.put(word, Math.log((double) totalDocs / (docCount + 1)));
        }

        return idfMap;
    }
}